from dotenv import load_dotenv
from langchain.retrievers import MultiQueryRetriever

from langchain_chroma import Chroma

from langchain_huggingface import HuggingFaceEmbeddings
from langchain_openai import ChatOpenAI

load_dotenv()

embeddings = HuggingFaceEmbeddings(
    model_name="D:/models/BAAIbge-base-zh-v1.5",
    model_kwargs={'device': 'cpu'}
)

#llm = ChatOpenAI(model="xdeepseekv32exp", temperature=0.7)
llm = ChatOpenAI(model="deepseek-ai/DeepSeek-V3.2-Exp", temperature=0.7)

vectorstore = Chroma(
    persist_directory="./chroma_db_rag_3",
    embedding_function=embeddings
)

# 从LLM 和向量数据库创建一个MultiQueryTetriever  它会自动处理 生成查询 ->  检索 -> 合并去重的整个流程
multiquery_retriever = MultiQueryRetriever.from_llm(
    retriever=vectorstore.as_retriever(),  # 使用创建的向量数据库作为基础检索器
    llm=llm  # 使用大模型来生成子查询
)

# 使用原始查询进行调用
user_query = "明朝有几个皇帝？"
print(llm.invoke("你好,测试一下"))
retrieved_docs = multiquery_retriever.invoke(user_query)

print(retrieved_docs, f"扩展更多问题(MultiQuery) 对 {user_query} 的检索结果")

# 深入了解它生成了哪些子查询
import logging

logging.basicConfig()
logging.getLogger("langchain.retriever.multi_query_retriever").setLevel(logging.INFO)
retrieved_docs = multiquery_retriever.invoke(user_query)
